import asyncio
import os
from datetime import datetime

from llama_index.core import SimpleDirectoryReader, Document

from base_rag import RAG
from utils import ocr_file_to_text_llm

class TraditionalRAG(RAG):
    async def load_data(self):
        docs = []
        for file in self.files:
            # 对图片及文档进行OCR识别
            contents = ocr_file_to_text_llm(file)

            temp_file = datetime.now().strftime("%Y%m%d%H%M%S") + ".txt"
            with open(temp_file, "w", encoding="utf-8") as f:
                f.write(contents)

            data = SimpleDirectoryReader(input_files=[temp_file]).load_data()
            doc = Document(text="\n\n".join([d.text for d in data[0:]]), metadata={"path": file})
            docs.append(doc)
            os.remove(temp_file)
        return docs

if __name__ == "__main__":
    rag = TraditionalRAG(files=["../test_data/222.jpg"])
    asyncio.run(rag.create_index_local())
